package profiler;

import org.objectweb.asm.ClassAdapter;
import org.objectweb.asm.ClassVisitor;
import org.objectweb.asm.Label;
import org.objectweb.asm.MethodVisitor;
import org.objectweb.asm.Opcodes;
import org.objectweb.asm.Type;

public class Instrumentor extends ClassAdapter implements Opcodes {

    private final Filter filter;

    private boolean accepted;

    private String className;

    public Instrumentor(ClassVisitor arg0, Filter filter) {
        super(arg0);
        this.filter = filter;
    }

    @Override
    public void visit(int arg0, int arg1, String arg2, String arg3, String arg4, String[] arg5) {
        accepted = filter.instrumentClass(arg2, arg1);
        if (accepted) {
            System.out.println("instrumenting: " + arg2);
        }
        className = arg2;
        super.visit(arg0, arg1, arg2, arg3, arg4, arg5);
    }

    @Override
    public MethodVisitor visitMethod(int arg0, String arg1, String arg2, String arg3, String[] arg4) {
        if (accepted && filter.includeMethod(className, arg0, arg1, arg2, arg3, arg4)) {
            String newName = AgentBootstrap.INSTRUMENTATION_PREFFIX + arg1;
            System.out.println("creating stub method: " + newName);
            createStubMethod(newName, arg0, arg1, arg2, arg3, arg4);
            return super.visitMethod(arg0, newName, arg2, arg3, arg4);
        } else {
            return super.visitMethod(arg0, arg1, arg2, arg3, arg4);
        }
    }

    private void createStubMethod(String newName, int modifiers, String methodName, String methodDescriptor, String signature,
            String[] exceptions) {
        MethodVisitor mv = super.visitMethod(modifiers, methodName, methodDescriptor, signature, exceptions);
        Label l0 = new Label();
        Label l1 = new Label();
        Label l2 = new Label();

        Type returnType = getReturnTypeForDescriptor(methodDescriptor);
        mv.visitTryCatchBlock(l0, l1, l2, null);
        mv.visitMethodInsn(INVOKESTATIC, "profiler/SingletonTimer", "getInstance", "()Lprofiler/Timing;");
        mv.visitLdcInsn(className + "," + methodName + "," + methodDescriptor);
        mv.visitMethodInsn(INVOKEINTERFACE, "profiler/Timing", "startMethodCall", "(Ljava/lang/String;)V");
        mv.visitLabel(l0);

        doMethodInvocation(newName, modifiers, methodDescriptor, mv, returnType);

        mv.visitLabel(l1);
        mv.visitMethodInsn(INVOKESTATIC, "profiler/SingletonTimer", "getInstance", "()Lprofiler/Timing;");
        mv.visitMethodInsn(INVOKEINTERFACE, "profiler/Timing", "stop", "()V");

        if (!returnType.equals(Type.VOID_TYPE)) {
            mv.visitVarInsn(returnType.getOpcode(ILOAD), 2);
            mv.visitInsn(returnType.getOpcode(IRETURN));
        } else {
            mv.visitInsn(RETURN);
        }

        mv.visitLabel(l2);
        mv.visitLineNumber(23, l2);
        mv.visitFrame(Opcodes.F_SAME1, 0, null, 1, new Object[]{ "java/lang/Throwable" });
        mv.visitVarInsn(ASTORE, 1);
        mv.visitMethodInsn(INVOKESTATIC, "profiler/SingletonTimer", "getInstance", "()Lprofiler/Timing;");
        mv.visitMethodInsn(INVOKEINTERFACE, "profiler/Timing", "stop", "()V");
        mv.visitVarInsn(ALOAD, 1);
        mv.visitInsn(ATHROW);
        mv.visitMaxs(2, 3);
        mv.visitEnd();
    }

    private void doMethodInvocation(String newName, int modifiers, String methodDescriptor, MethodVisitor mv, Type returnType) {
        final int opcode;
        final int localArgsOffset;
        if (isStatic(modifiers)) {
            localArgsOffset = 0;
            opcode = INVOKESTATIC;
        } else {
            mv.visitVarInsn(ALOAD, 0);
            localArgsOffset = 1;
            opcode = INVOKESPECIAL;
        }

        loadMethodParams(mv, getParameterTypeForDescriptor(methodDescriptor), localArgsOffset);
        mv.visitMethodInsn(opcode, className, newName, methodDescriptor);
        if (!returnType.equals(Type.VOID_TYPE)) {
            mv.visitVarInsn(returnType.getOpcode(ISTORE), 2);
        }
    }

    private void loadMethodParams(MethodVisitor mv, Type[] parameterTypeForDescriptor, int localArgsOffset) {
        int offset = localArgsOffset;
        for (int i = 0; i < parameterTypeForDescriptor.length; i++) {
            Type type = parameterTypeForDescriptor[i];
            mv.visitVarInsn(type.getOpcode(ILOAD), localArgsOffset++);
            if (type.equals(Type.LONG_TYPE) || type.equals(Type.DOUBLE_TYPE)) {
                offset++;
            }
        }
    }

    private static Type getReturnTypeForDescriptor(String descriptor) {
        return Type.getReturnType(descriptor);
    }

    private static Type[] getParameterTypeForDescriptor(String descriptor) {
        return Type.getArgumentTypes(descriptor);
    }

    private static boolean isStatic(int modifiers) {
        return (Opcodes.ACC_STATIC & modifiers) != 0;
    }

    /**
     * @return the accepted
     */
    public boolean isAccepted() {
        return accepted;
    }

}
